#!/usr/bin/env python
"""
Simple command line tool to extract sequences of interest from a sequence file.

Description:

Program to reformat or filter sequences in a variety of formats.
See http://www.biopython.org/wiki/SeqIO for a list of formats,
although there are currently more supported by biopython.

e.g. `sff` format (Roche 454's binary sequence format) is supported
by biopython, but not mentioned on the webpage.

e.g. ::

    alignDB.py  -db infile.sto
                [-format stockholm]
                [-accs "acc_1" "acc_2" "acc_N"]
                [-desc "Methanobacter maripuladis" "Methanococcus"]

"""

# Copyright (c) Alex Leach. 2011 - 2013
# Contact: Alex Leach (albl500@york.ac.uk)
# University of York, Department of Biology,
# Wentworth Way, York YO10 4DU, United Kingdom
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have receive a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>

from ssummo.cmd_options import Options

import re
import sys

from Bio import AlignIO

class LocalOptions(Options):

    options = { '-db' : None ,
                '-accs': [] ,
                '-desc' : [] ,
                '-format' : 'fasta',
                '-out' : sys.stdout,
                '-outformat' : 'fasta',
                '--reverse' : False
            }

    help_text = { '-db' : 'File which contains sequences of interest',
                '-accs' : 'List of accessions to look for within <-db>. This is '+
                          'relaxed, so returns all matches which contain a given '+
                          'accession, even if it isnt the complete accession. '+
                          'Separate with spaces.',
                '-desc' : 'Descriptions to look for within db. Again, is relaxed, '+
                          'so returns all matches. Envelope each description with '+
                          '" characters and separates with spaces.',
                '-format' : 'The input sequence format. Default: fasta',
                '-out' : 'Output file. Default: standard out. N.B. Appends to the '+
                         'file.',
                '-outformat' : 'Output file format. Default: fasta.',
                '--reverse' : 'Reverse transcribe the sequences'
              }

    def print_help(self):
        print('Valid options:-')
        for key, val in self.help_text.items():
            print( '{0} : {1}'.format(key, val).rjust(80) )

    def parse_args(self, args):
        if len(args) == 0:
            self.print_help()
            exit()
        options = self.options
        prevargDash = False
        for i in xrange(len(args)):
            if args[i].startswith('--'):
                other = options[args[i]]
                options[args[i]] = not other
            elif args[i].startswith('-'):
                if args[i] in options.keys():
                    prevargDash = True
                    command = args[i]
                else:
                    print( 'Invalid option: {0}'.format(args[i]) )
                    self.print_help()
                    exit()
            elif prevargDash:
                if command in ['-db', '-outformat', '-out']:
                    options.update( { command :  args[i] } )
                    prevargDash = False
                elif command == '-accs':
                    options[command].append(args[i])
                elif command == '-desc':
                    options[command].append( '('+args[i] + ')')
                elif command == '-format':
                    options.update( {command : args[i] } )
                    prevargDash = False
            else:
                print("Don't know what to do with {0}".format(args[i]))
                exit()
        return options

class Reader():

    def __init__(self, options):
        if len(options['-desc']) > 0:
            self.desc_reg = re.compile( '|'.join( options['-desc'] ), re.I )
        if len(options['-accs']) > 0:
            self.acc_reg = re.compile( '|'.join( options['-accs'] ), re.I )
        if options['--reverse']:
            self._yield = self.reverser
        else:
            self._yield = self.writer
        self.options = options

    def reverser(self, seq):
        seq.seq = seq.seq.reverse_complement()
        return seq

    def writer(self, seq):
        return seq

    def __iter__(self):
        Ndesc = len(options['-desc'])
        Naccs = len(options['-accs'])
        count = 0
        if Naccs == 0 and Ndesc == 0:
            printed = False
        else:
            printed = True
        for seq in AlignIO.parse(options['-db'], options['-format']):
            count += 1
            if Naccs > 0:
                accFind = self.acc_reg.search( seq.id )
                if accFind:
                    yield self._yield( seq )
                    printed = True
            elif Ndesc > 0:
                desc_find = self.desc_reg.search(seq.description)
                if desc_find:
                    yield self._yield(seq)
                    printed = True
            else:
                yield self._yield( seq )
        if not printed and count > 0:
                sys.stderr.write( "No matching sequences found. Here's the last sequence header as an example...\n" )
                sys.stderr.write( seq.format( options['-outformat']).splitlines()[0]  )
                sys.stderr.write('\n')
        return

def parse_input(options):
    aligner = Reader( options )
    if options['-out'] != sys.stdout:
        out = file( options['-out'] , 'w' )
    else:
        out = sys.stdout
#    print sequences[0],type(sequences[0])
#    alignment_object = MultipleSeqAlignment( sequences )
    AlignIO.write( aligner , out , options['-outformat'] )

if __name__ == '__main__':
    args = sys.argv[1:]
    options = LocalOptions().parse_args(args)
    parse_input(options)
